Skip to content

Feat/schema improvements and optimization - #95

Open
carlosengutierrez wants to merge 9 commits into
mainfrom
feat/schema-improvements
Open

Feat/schema improvements and optimization#95
carlosengutierrez wants to merge 9 commits into
mainfrom
feat/schema-improvements

Conversation

@carlosengutierrez

Copy link
Copy Markdown
Collaborator

Adds parameter optimization to Neuro-Workflow: a search loop that runs an existing workflow
repeatedly, tunes declared parameters and reports how close each run lands to a declared target.

A workflow is a DAG that runs once; a search is a loop that runs it many times, so the loop sits
outside the graph rather than becoming a node inside it.

What is here

  • src/neuroworkflow/optimization/ — the engine: dotted addressing (Node.parameter[.key] to
    tune, Node.output_port[.key] to measure), the spec, a file-based ledger, and optimizer backends
    (random plus Optuna's CMA-ES, TPE, NSGA-II/III). Single-objective backends refuse a
    multi-objective spec rather than inventing weights between objectives.
  • NW_Optimization node — declares how to search: algorithm, budget, seed. No ports and no
    process steps, so a workflow containing it executes exactly as it would without it. Its presence
    on the canvas is the signal for the generator to emit an optimization run instead of a single
    execution. Replaces JointOptimizationNode, which was an in-graph grid search.
  • Signature-based network reuseNW_SimConfig rebuilds the SONATA network only when
    something structural changed, so a search does not rebuild it once per trial. Nodes opt in with
    REUSABLE_PATHS; the engine itself stays domain-free.
  • Measurements on NW_Analysisfiring_rate_hz and isi_stats per population, so a target
    has something to be compared against.
  • Schemaunit and measures on ParameterDefinition; optimization_range accepts a
    per-key dict for dict-valued parameters. Each node instance now carries its own
    NODE_DEFINITION copy, so marking exc optimizable no longer affects inh.

Reporting across objectives in different units

A miss in Hz and a miss in ms cannot be added. Each objective reports its own miss in its own unit,
and the single comparable figure is the worst objective's miss divided by the width of its target
range — printed as furthest from target: probe_isi, 1.4x its target range and stored as
target_ranges_off. Pareto dominance is scale-free and NSGA-II rescales internally, so this
affects reporting and ranking only; what the optimizer is told stays raw and per-objective.

Docs and examples

docs/OPTIMIZATION.md for how the engine works, docs/OPTIMIZATION_GUI_HANDOFF.md for the GUI
half. Five notebooks under notebooks/, two of which show what the code generator should emit.

Not done, and needed before the GUI can use this

optuna and cmaes are not in the nest kernel image. NW_Optimization defaults to cmaes,
so the first optimization generated in the GUI fails on import until one line is added at
Dockerfile.nest:128. Left out of this PR because it forces a rebuild of the nest image; it is
written up as step 1 of the handoff document.

Testing

Single- and multi-objective searches run end to end in Jupyter against real NEST/BMTK, and with toy
nodes for the engine paths. GUI copies under codes/ are synced and byte-identical to src/.

@IzBrain67

Copy link
Copy Markdown
Collaborator

Thank you for this — the overall direction is right. Putting the search loop outside the workflow, the file-based ledger that a person, the GUI, or an agent can all read and steer, and the per-population firing-rate / ISI outputs on the analysis node are all things users will benefit from, optimization or not. The reasoning captured in the docs is also a real asset for whoever touches this next.

Before merging I'd like to settle one design question and fix a handful of things that would bite users on first use. I'm listing them symptom-first. I also have a detailed, file-by-file list and will hand it over in stages so this doesn't have to be fixed all at once.

One design question to settle first

The PR currently describes three different homes for "what to explore and what to hit":

  • the NW_Optimization node's own docstring says it deliberately does not hold them (they live on each parameter),
  • the GUI handoff document says the study — all three parts — should live in the NW_Optimization node, keyed by node id,
  • the skill document recommends declaring the target on the study (add_objective) and argues against adding a parameter the simulation never reads — yet the same PR adds a mean_firing_rate parameter to NW_Population that does exactly that.

The GUI work depends heavily on which of these is the intended model, so I'd rather pick one now than build the panel on the wrong one. My preference is the handoff-document version (study lives in the optimization node, referencing nodes by id), and dropping mean_firing_rate from NW_Population. Happy to discuss.

Things that would bite users on first use

  1. A measurement that comes back as NaN counts as "target reached". If any node produces a NaN (not a number), the run reports that the target was met and stops. NaN also shows up in the list of measurable values as if it were a valid target. NaN should be treated as a failed trial.

  2. The default algorithm (CMA-ES) silently throws away most of the simulations. With pop_size=16, only the first 6 trials of each generation are used to learn; the other 10 are run but ignored. I verified this against Optuna 5.0: the sampler needs to be told the population size (it's already done for NSGA, just not for CMA-ES). One-line fix, but it changes results a lot.

  3. Steering a run with inject can crash it. If a candidate in control.json lists only some of the parameters (which is the natural thing to write), the run dies with an error after evaluating it, and status.json stays at "running" forever. Injected candidates should be validated and completed with the current values.

  4. A node that doesn't re-emit an output is scored with the previous trial's value. The engine doesn't clear output ports between trials, and a node that succeeds without producing an output only prints a warning. The shipped analysis node always re-emits, but any custom node could be silently scored on stale data.

  5. Connection rules that aren't plain lambdas now crash network setup. The new "rebuild only when something changed" check assumes every callable has function internals; a functools.partial or a callable object (which BMTK accepted before) fails with an error during setup.

  6. Connection-rule strings from the GUI are evaluated as arbitrary Python. A shared workflow could contain any expression, not just a lambda, and it runs in whoever opens it. At minimum this should be restricted to lambda expressions and documented; longer term, a whitelist of rule functions.

  7. The first optimization generated in the GUI fails on import. As you noted, optuna/cmaes aren't in the nest kernel image and the node defaults to cmaes. Let's include the Dockerfile line and a pip install -e ".[optimization]" extra in this PR or an immediate follow-up so the feature is usable end to end.

Smaller things worth fixing in the same pass

  • Changing a global dict/list/array that a connection rule reads does not trigger a rebuild, so the old network is reused silently in the default auto mode.
  • Networks built with a random connection rule are now frozen after the first run (previously each re-run drew a new realisation). Fine for optimization, but it's a behaviour change for normal re-runs and the overwrite parameter description no longer matches.
  • Numpy scalar outputs (np.int64, np.float32) are not recognised as numbers, so they vanish from the measurable list; only Python floats and np.float64 work.
  • Injected candidates are recorded as "source": "ask" in the ledger, contrary to the docs.
  • Defaults disagree: AlgorithmConfig defaults to random, the node to cmaes; the max_generations help text promises early stopping that doesn't happen for multi-objective runs.

Housekeeping

  • examples/neuron_optimization.py still imports the deleted JointOptimizationNode and no longer runs.
  • The new optimization/ files don't pass black / isort (repo convention in CLAUDE.md); a quick format pass fixes it.
  • There are no automated tests. The engine runs without Optuna using the random algorithm, so a small pytest suite with toy nodes (the ones you used for the engine paths) would be very welcome — it's the one thing that protects this code through the next round of changes.
  • The notebooks add ~57k lines, mostly saved outputs (~2.8 MB). Could we strip outputs, or keep only the two generated examples with outputs?

Let's hold off on merging until the design question above is settled — thanks again for the careful work here.

Reviewed with the help of Claude Code.

Treat NaN as a failed trial, complete injected candidates, pass CMA-ES popsize,
restrict connection_rule strings to lambdas, and default NW_Optimization to
random so a search runs without Optuna. Add toy-node tests, strip notebook
outputs, and document that GUI generate-code and the nest image extra remain
follow-ups.
@kirillmitrofanov

Copy link
Copy Markdown
Collaborator

Audit on this branch (engine harden, not a rewrite): the outer-loop search is applicable and reasonable. I pushed first-use fixes: NaN is a failed trial (not “in target”), inject completes missing dims and records source: inject, CMA-ES gets popsize, output ports are cleared between trials, _stable no longer crashes on functools.partial, and connection_rule strings must be a lambda (no full-eval builtins). NW_Optimization now defaults to random so a run works without Optuna; pip install -e ".[optimization]" is in pyproject.toml. Toy-node pytest is in tests/test_optimization_engine.py; notebook outputs were stripped.

Open question for @carlosengutierrez (please do not merge until this is answered): where does the study live — per-parameter optimizable/is_objective (as shipped, including mean_firing_rate on NW_Population), vs the GUI-handoff model (study on NW_Optimization keyed by node id), vs spec.add_objective() in the skill? @IzBrain67 prefers the handoff version and dropping mean_firing_rate. This pass did not pick.

GUI generate-code and adding optuna/cmaes to Dockerfile.nest remain follow-ups (needs an ops image rebuild). Live dbrain.jp is still the old bundle — palette search for “Optimization” does not show NW_Optimization. Merge later with #92 will need a hand merge on schema.py / node.py.

@IzBrain67

Copy link
Copy Markdown
Collaborator

@carlosengutierrez — one decision is needed from you before this can move, and it is a product question, not a code one.

Where should a user set up an optimization?

  • A (recommended): everything about a search — which parameters to explore, the target to hit, and the algorithm — is set in one place, on the Optimization node. Users see the whole study in one panel, can keep two studies over the same model, and
    the neuron nodes stay clean.
  • B (as currently shipped): the search settings are spread across the individual nodes (each parameter carries its own "optimizable" and "target" flags), and the Optimization node only holds the algorithm.

Both work with the engine you wrote; the difference is what users will see in the GUI, and the GUI panel has to be built for one of them.

Kirill has already fixed the correctness items from my review, so this is the only thing holding the PR. If I don't hear back by Tuesday, September 16, I'll go ahead with A and adjust the docs and the mean_firing_rate parameter to match. A
one-line reply ("A" or "B") is enough.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants